import numpy as np
from parl.utils import ReplayMemory


class Dim2DReplayMemory(ReplayMemory):

    def __init__(self, max_size, obs_dim, act_dim):
        """ create a replay memory for off-policy RL or offline RL.

        Args:
            max_size (int): max size of replay memory
            obs_dim (list or tuple): observation shape
            act_dim (list or tuple): action shape
        """
        self.max_size = int(max_size)
        self.obs_dim = obs_dim
        self.act_dim = act_dim

        self.obs = np.zeros((max_size, *obs_dim), dtype='float32')
        if act_dim == 0:  # Discrete control environment
            self.action = np.zeros((max_size,), dtype='int32')
        else:  # Continuous control environment
            self.action = np.zeros((max_size, act_dim), dtype='float32')
        self.reward = np.zeros((max_size,), dtype='float32')
        self.terminal = np.zeros((max_size,), dtype='bool')
        self.next_obs = np.zeros((max_size, *obs_dim), dtype='float32')

        self._curr_size = 0
        self._curr_pos = 0
